# -*- coding: utf-8 -*-
# =============================================================================
#     FileName:
#         Desc:
#       Author:
#        Email:
#     HomePage:
#      Version:
#   LastChange:
#      History:
# =============================================================================
import sys
import time
import decimal

import fire
import redis
import argparse
import logging
from collections import namedtuple
from utils.logger_helper import LoggerHelper
from utils.top_max_heap import TopMaxHeapItem, TopMaxHeap
from utils.human_helper import HumanHelper

APP_NAME = "redis_big_keys"
APP_VERSION = "1.0"
logger = LoggerHelper.get_logger()

RedisConfig = namedtuple("RedisConfig", ["redis_host", "redis_port", "redis_password", "redis_db"])


class BigKeyScanner(object):
    def __init__(self, redis_config: RedisConfig,
                 min_key_size=1 * 1024, big_key_count=100, big_key_type="all",
                 total_scan_count=0, batch_scan_count=100, batch_scan_interval: float = 0.01):
        self.redis_config = redis_config
        self.min_key_size = min_key_size
        self.big_key_count = big_key_count
        self.big_key_type = big_key_type
        self.top_heap = TopMaxHeap(top_size=big_key_count)
        self.total_scan_count = total_scan_count
        self.batch_scan_count = batch_scan_count
        self.batch_scan_interval = batch_scan_interval
        self.scan_key_index = 0

    def get_redis_conn(self):
        return redis.StrictRedis(
            host=self.redis_config.redis_host,
            port=self.redis_config.redis_port,
            password=self.redis_config.redis_password,
            db=self.redis_config.redis_db
        )

    @classmethod
    def get_key_info_by_string(cls, key_info_str: str):
        key_info = {}
        sub_items = key_info_str.split(" ")
        for sub_item in sub_items:
            if sub_item.find(":") > 0:
                tmp_list = sub_item.split(":")
                key_info[tmp_list[0]] = tmp_list[1]
        return key_info

    def get_batch_key_info(self, key_name_list):
        redis_conn = self.get_redis_conn()
        key_info_list = list()
        with redis_conn.pipeline(transaction=False) as pipe:
            for key_name in key_name_list:
                pipe.memory_usage(key=key_name)
            memory_list = pipe.execute()
            for key_name in key_name_list:
                pipe.type(name=key_name)
            type_list = pipe.execute()
        for key_index in range(0, len(key_name_list)):
            key_memory = memory_list[key_index]
            key_type = type_list[key_index].decode("utf-8")
            key_name = key_name_list[key_index].decode("utf-8")
            if (key_memory is None) or (key_type is None):
                continue
            key_info = {
                "key_name": key_name,
                "key_type": key_type,
                "key_memory": key_memory,
            }
            key_info_list.append(key_info)
        return key_info_list

    def scan_redis_keys(self):
        logger.info("开始扫描")
        redis_conn = self.get_redis_conn()
        start_cursor = 0
        while True:
            if self.total_scan_count > 0:
                if self.scan_key_index > self.total_scan_count:
                    break
            next_cursor, key_name_list = redis_conn.scan(cursor=start_cursor, match="*", count=self.batch_scan_count)
            if next_cursor != 0:
                start_cursor = next_cursor
            else:
                break
            key_info_list = self.get_batch_key_info(key_name_list=key_name_list)
            for key_info in key_info_list:
                key_memory = key_info["key_memory"]
                if key_memory < self.min_key_size:
                    continue
                top_item = TopMaxHeapItem(item_key=key_memory, item_data=key_info)
                self.top_heap.push_item(top_item)
            self.scan_key_index += self.batch_scan_count
            if self.scan_key_index % 10000 == 0:
                logger.info("完成第{}个KEY扫描".format(self.scan_key_index))
            time.sleep(self.batch_scan_interval)
        logger.info("完成扫描")
        self.show_big_keys()

    def show_big_keys(self):
        top_items = self.top_heap.get_top_items()
        logger.info("共获取到{}个大KEY(超过{})".format(len(top_items), self.min_key_size))
        top_messages = list()
        top_messages.append("")
        top_message_demo = "{:>15} {:>10}     {} "
        top_messages.append(top_message_demo.format("key_memory", "key_type", "key_name"))
        for top_item in top_items:
            key_info = top_item.item_data
            top_messages.append(top_message_demo.format(
                HumanHelper.format_byte_number(key_info["key_memory"]),
                key_info["key_type"],
                key_info["key_name"]
            ))
        logger.info("\n".join(top_messages))
        pass


class CommandParser(object):

    @classmethod
    def get_parser(cls):
        parser = argparse.ArgumentParser(
            description='从Redis扫描BIG KEY',
            add_help=False)
        connect_setting = parser.add_argument_group('参数信息')
        parser.add_argument(
            '--help',
            dest='help',
            action='store_true',
            help='获取帮助信息',
            default=False
        )
        connect_setting.add_argument(
            '--redis_host',
            dest='redis_host',
            type=str,
            default="",
            help='Redis服务器地址,默认值为127.0.0.1'
        )
        connect_setting.add_argument(
            '--redis_port',
            dest='redis_port',
            type=int,
            default=6379,
            help='Redis服务器端口,默认值为6379'
        )
        connect_setting.add_argument(
            '--redis_db',
            dest='redis_db',
            type=int,
            default=0,
            help="Redis服务器DB,默认值为0"
        )
        connect_setting.add_argument(
            '--redis_password',
            dest='redis_password',
            type=str,
            default="",
            help="Redis密码,默认值为空字符串"
        )
        connect_setting.add_argument(
            '--min_key_size',
            dest='min_key_size',
            type=int,
            default=1000,
            help="最小key大小(单位Byte),默认值1000"
        )
        connect_setting.add_argument(
            '--top_key_size',
            dest='top_key_size',
            type=int,
            default=100,
            help="返回大Key的数量,默认值100"
        )
        connect_setting.add_argument(
            '--top_key_type',
            dest='top_key_type',
            type=str,
            default="all",
            help="要扫描的大KEY类型，默认值为ALL"
        )
        connect_setting.add_argument(
            '--big_key_count',
            dest='big_key_count',
            type=int,
            default=0,
            help="扫描KEY的数量，默认值0(扫描全部),当扫描KEY数量大于0时，仅扫描指定数量的key便停止扫描"
        )
        connect_setting.add_argument(
            '--batch_scan_interval',
            dest='batch_scan_interval',
            type=float,
            default=0.01,
            help="每次扫描的休眠间隔,默认每次扫描休眠0.01秒"
        )
        connect_setting.add_argument(
            '--batch_scan_count',
            dest='batch_scan_count',
            type=int,
            default=100,
            help="每次扫描的大小,默认每次扫描100个Key"
        )
        return parser

    @classmethod
    def parse_args(cls, command_args):
        need_print_help = False if command_args else True
        parser = cls.get_parser()
        args = parser.parse_args(command_args)
        if args.help or need_print_help:
            parser.print_help()
            sys.exit(1)
        return args


def init_logger(with_debug: bool = False):
    if with_debug:
        LoggerHelper.init_logger(logger_level=logging.DEBUG)
    else:
        LoggerHelper.init_logger(logger_level=logging.INFO)
    global logger
    logger = LoggerHelper.get_logger()


def main(redis_port: int = 6379, redis_host: str = "127.0.0.1",
         redis_password: str = "", redis_db: int = 0,
         min_key_size: int = 1024, big_key_type: str = "all",
         big_key_count: int = 100, total_scan_count: int = 0,
         batch_scan_count: int = 100, batch_scan_interval: float = 0.01,
         with_debug: bool = False):
    """

    :param redis_host: Redis实例地址
    :param redis_port: Redis实例端口
    :param redis_password: Redis实例账号
    :param redis_db: Redis服务器DB
    :param min_key_size: 最小key大小(单位Byte)
    :param big_key_type: 要扫描的大KEY类型
    :param big_key_count: 要返回大Key的数量
    :param total_scan_count: 要扫描KEY的数量, 默认值0(扫描全部),当扫描KEY数量大于0时，仅扫描指定数量的key便停止扫描
    :param batch_scan_count: 每次扫描的KEY的数量, 单次扫描过多会影响Redis实例的稳定性。
    :param batch_scan_interval: 每次扫描的休眠间隔,单位为秒
    :param with_debug: 是否开启DEBUG模式
    :return:
    """
    init_logger(with_debug=with_debug)
    bks = BigKeyScanner(
        redis_config=RedisConfig(
            redis_host=redis_host,
            redis_port=redis_port,
            redis_password=redis_password,
            redis_db=redis_db
        ),
        min_key_size=min_key_size,
        big_key_count=big_key_count,
        big_key_type=big_key_type,
        total_scan_count=total_scan_count,
        batch_scan_interval=batch_scan_interval,
        batch_scan_count=batch_scan_count
    )
    bks.scan_redis_keys()


if __name__ == '__main__':
    fire.Fire(main)
